from flask import Blueprint, current_app, jsonify,request,g
from flask import send_file
from selenium import webdriver
import selenium
from io import BytesIO
import base64
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.options import Options
from PIL import Image
import time

print(selenium.__version__)

# 创建 ChromeOptions 对象
chrome_options = Options()

# 设置 Chrome 为无头模式
chrome_options.add_argument('--headless')
# chrome_options.add_argument("device-scale-factor=0.5")



curveprint = Blueprint('curveprint', __name__)

@curveprint.route('/api/screenshot',methods=['GET','POST'])
def screenshot():
    w = int(request.args.get('width',800))
    h = int(request.args.get('height',600))
    degrees = int(request.args.get('degrees',0))
    
    # 创建 Chrome 浏览器对象
    driver = webdriver.Chrome(options=chrome_options)

    driver.set_window_size(w-20,h-50)

    # driver.get("about:blank")

    # 打开指定网页
    driver.get(changedUrl(request.url))

    driver.set_window_size(w,h)
    
    # 调整窗口大小以避免滚动条
    # set_window_size_to_avoid_scroll(driver, w, h)

    # driver.execute_script("document.body.style.zoom='100%'")

    # driver.get('http://www.baidu.com')

    
    wait = WebDriverWait(driver, 10)

    # 等待网页加载完成

    # 等待网页标题包含指定字符串
    wait.until(EC.title_contains('Example'))

    # 等待网页标题包含指定字符串
    # wait.until(EC.title_contains('complete'))

    # 截取网页截图
    screenshot = driver.get_screenshot_as_png()

    
    # 将截图保存到本地
    # with open('files/screenshot.png', 'wb') as f:
    #     f.write(screenshot)

    # 将截图转换为 base64 编码
    encoded_screenshot = base64.b64encode(screenshot).decode('utf-8')

    # 关闭浏览器对象
    driver.quit()

    
    # 将截图数据加载到 Image 对象中
    image = Image.open(BytesIO(screenshot))

    # 旋转图像 90 度
    rotated_image = image.rotate(degrees, expand=True)
    

    # 将旋转后的图像保存到内存中
    output_buffer = BytesIO()
    rotated_image.save(output_buffer, format="PNG")
    output_buffer.seek(0)  # 将缓冲区的读取位置设置为开头

    # 返回原始二进制图像数据
    return send_file(output_buffer, mimetype='image/png', as_attachment=False)

def changedUrl(url):
    base_url='http://172.16.14.99/curve/prt?' + url.split('?')[1]
    # base_url='http://127.0.0.1:3000/curve/prt?' + url.split('?')[1]
    return base_url

def set_window_size_to_avoid_scroll(driver, width, height, max_iterations=20, increment=10):
    iteration = 0
    while iteration < max_iterations:
        driver.set_window_size(width, height)
        
        # 检查是否有垂直滚动条
        has_vertical_scroll = driver.execute_script("return document.documentElement.scrollHeight > document.documentElement.clientHeight;")
        # 检查是否有水平滚动条
        has_horizontal_scroll = driver.execute_script("return document.documentElement.scrollWidth > document.documentElement.clientWidth;")

        if not has_vertical_scroll and not has_horizontal_scroll:
            # print(width,height)
            break

        if has_vertical_scroll:
            height += increment

        if has_horizontal_scroll:
            width += increment

        iteration += 1
